import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.stattools import coint
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import seaborn as sns
# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体显示中文
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
# 初始化tushare pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')

def get_stock_data(ts_code, start_date, end_date):
    """获取股票日线数据"""
    df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df = df.sort_values('trade_date').set_index('trade_date')
    return df

def get_index_data(index_code, start_date, end_date):
    """获取指数日线数据"""
    df = pro.index_daily(ts_code=index_code, start_date=start_date, end_date=end_date)
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df = df.sort_values('trade_date').set_index('trade_date')
    return df

# 获取示例数据
start_date = '20200101'
end_date = '20231231'

# 获取个股数据
stock_a = get_stock_data('601318.SH', start_date, end_date)  # 中国平安
stock_b = get_stock_data('600036.SH', start_date, end_date)  # 招商银行

# 获取指数数据
hs300 = get_index_data('000300.SH', start_date, end_date)    # 沪深300




#2. 量价策略 - 双均线策略
def moving_average_strategy(data, short_window=5, long_window=20):
    """双均线策略"""
    signals = pd.DataFrame(index=data.index)
    signals['price'] = data['close']
    signals['short_ma'] = data['close'].rolling(window=short_window, min_periods=1).mean()
    signals['long_ma'] = data['close'].rolling(window=long_window, min_periods=1).mean()

    # 生成交易信号
    signals['signal'] = 0.0
    signals['signal'][short_window:] = np.where(
        signals['short_ma'][short_window:] > signals['long_ma'][short_window:], 1.0, 0.0)

    # 计算实际持仓变化
    signals['positions'] = signals['signal'].diff()

    return signals


# 应用双均线策略
ma_signals = moving_average_strategy(stock_a)

# 可视化
plt.figure(figsize=(12, 6))
plt.plot(ma_signals['price'], label='Price')
plt.plot(ma_signals['short_ma'], label=f'{5}日均线')
plt.plot(ma_signals['long_ma'], label=f'{20}日均线')

# 标记买入信号
plt.plot(ma_signals.loc[ma_signals['positions'] == 1.0].index,
         ma_signals['short_ma'][ma_signals['positions'] == 1.0],
         '^', markersize=10, color='g', label='买入')

# 标记卖出信号
plt.plot(ma_signals.loc[ma_signals['positions'] == -1.0].index,
         ma_signals['short_ma'][ma_signals['positions'] == -1.0],
         'v', markersize=10, color='r', label='卖出')

plt.legend()
plt.title('双均线交易策略')
plt.show()


#3. K线形态策略 - 锤子线和吞没形态
def candlestick_patterns(data):
    """识别K线形态"""
    patterns = pd.DataFrame(index=data.index)

    # 锤子线识别
    body = abs(data['open'] - data['close'])
    upper_shadow = data['high'] - data[['open', 'close']].max(axis=1)
    lower_shadow = data[['open', 'close']].min(axis=1) - data['low']

    # 锤子线条件: 小实体，长下影线，短上影线
    hammer = (body <= 0.3 * (data['high'] - data['low'])) & \
             (lower_shadow >= 2 * body) & \
             (upper_shadow <= 0.1 * (data['high'] - data['low']))

    # 看涨吞没形态
    bullish_engulfing = (data['close'] > data['open']) & \
                        (data['close'].shift(1) < data['open'].shift(1)) & \
                        (data['close'] > data['open'].shift(1)) & \
                        (data['open'] < data['close'].shift(1))

    patterns['hammer'] = hammer
    patterns['bullish_engulfing'] = bullish_engulfing

    return patterns


# 应用K线形态识别
patterns = candlestick_patterns(stock_a)

# 可视化示例
sample_data = stock_a.iloc[50:70].copy()
sample_patterns = candlestick_patterns(sample_data)

plt.figure(figsize=(12, 6))
for i in range(len(sample_data)):
    color = 'r' if sample_data['close'].iloc[i] < sample_data['open'].iloc[i] else 'g'
    plt.plot([i, i], [sample_data['low'].iloc[i], sample_data['high'].iloc[i]], color=color)
    plt.plot([i, i], [sample_data['open'].iloc[i], sample_data['close'].iloc[i]],
             linewidth=5, color=color)

    # 标记锤子线
    if sample_patterns['hammer'].iloc[i]:
        plt.text(i, sample_data['high'].iloc[i], 'H', ha='center', va='bottom', fontsize=12)

    # 标记吞没形态
    if sample_patterns['bullish_engulfing'].iloc[i]:
        plt.text(i, sample_data['high'].iloc[i], 'BE', ha='center', va='bottom', fontsize=12)

plt.xticks(range(len(sample_data)), sample_data.index.date, rotation=45)
plt.title('K线形态识别示例')
plt.show()

#4. 统计套利策略 - 配对交易
def find_cointegrated_pairs(price_data):
    """寻找协整的股票对"""
    n = price_data.shape[1]
    score_matrix = np.zeros((n, n))
    pvalue_matrix = np.ones((n, n))
    pairs = []

    for i in range(n):
        for j in range(i + 1, n):
            S1 = price_data.iloc[:, i]
            S2 = price_data.iloc[:, j]
            _, pvalue, _ = coint(S1, S2)
            pvalue_matrix[i, j] = pvalue
            if pvalue < 0.05:  # 显著性水平5%
                pairs.append((price_data.columns[i], price_data.columns[j], pvalue))

    return pvalue_matrix, pairs


def pairs_trading_strategy(stock1, stock2, lookback=20):
    """配对交易策略"""
    df = pd.DataFrame()
    df['stock1'] = stock1['close']
    df['stock2'] = stock2['close']

    # 计算价差
    df['spread'] = df['stock1'] - df['stock2']

    # 计算z-score
    df['mean'] = df['spread'].rolling(window=lookback).mean()
    df['std'] = df['spread'].rolling(window=lookback).std()
    df['zscore'] = (df['spread'] - df['mean']) / df['std']

    # 生成交易信号
    df['long_entry'] = df['zscore'] <= -2.0
    df['long_exit'] = df['zscore'] >= 0.0
    df['short_entry'] = df['zscore'] >= 2.0
    df['short_exit'] = df['zscore'] <= 0.0

    # 初始化仓位
    df['positions'] = 0

    # 生成交易信号
    long_position = False
    short_position = False

    for i in range(len(df)):
        if df['long_entry'].iloc[i] and not long_position:
            df.loc[df.index[i], 'positions'] = 1  # 买入stock1，卖出stock2
            long_position = True
        elif df['long_exit'].iloc[i] and long_position:
            df.loc[df.index[i], 'positions'] = 0
            long_position = False
        elif df['short_entry'].iloc[i] and not short_position:
            df.loc[df.index[i], 'positions'] = -1  # 卖出stock1，买入stock2
            short_position = True
        elif df['short_exit'].iloc[i] and short_position:
            df.loc[df.index[i], 'positions'] = 0
            short_position = False

    return df




# 寻找协整股票对
price_data = pd.concat([stock_a['close'], stock_b['close']], axis=1)
price_data.columns = ['stock_a', 'stock_b']
_, pairs = find_cointegrated_pairs(price_data)
from statsmodels.tsa.stattools import adfuller

def check_stationarity(data):
    """检查时间序列的平稳性"""
    result = adfuller(data)
    print(f"ADF Statistic: {result[0]}")
    print(f"p-value: {result[1]}")
    return result[1] < 0.05
stock_a_diff = stock_a['close'].diff().dropna()
stock_b_diff = stock_b['close'].diff().dropna()
price_data_diff = pd.concat([stock_a_diff, stock_b_diff], axis=1)
price_data_diff.columns = ['stock_a', 'stock_b']
_, pairs = find_cointegrated_pairs(price_data_diff)
print("找到的协整股票对:", pairs)

# 应用配对交易策略
if pairs:
    pt_signals = pairs_trading_strategy(stock_a, stock_b)

    # 可视化
    plt.figure(figsize=(14, 7))

    # 价差和z-score
    plt.subplot(211)
    plt.plot(pt_signals['spread'], label='价差')
    plt.axhline(pt_signals['mean'].iloc[-1], color='gray', linestyle='--')
    plt.axhline(pt_signals['mean'].iloc[-1] + 2 * pt_signals['std'].iloc[-1], color='r', linestyle='--')
    plt.axhline(pt_signals['mean'].iloc[-1] - 2 * pt_signals['std'].iloc[-1], color='g', linestyle='--')
    plt.legend()
    plt.title('价差和标准差通道')

    # 交易信号
    plt.subplot(212)
    plt.plot(pt_signals['zscore'], label='Z-score')
    plt.axhline(2, color='r', linestyle='--')
    plt.axhline(-2, color='g', linestyle='--')
    plt.axhline(0, color='gray', linestyle='--')

    # 标记交易信号
    plt.plot(pt_signals[pt_signals['positions'] == 1].index,
             pt_signals['zscore'][pt_signals['positions'] == 1],
             '^', markersize=10, color='g', label='买入信号')
    plt.plot(pt_signals[pt_signals['positions'] == -1].index,
             pt_signals['zscore'][pt_signals['positions'] == -1],
             'v', markersize=10, color='r', label='卖出信号')

    plt.legend()
    plt.title('Z-score和交易信号')
    plt.tight_layout()
    plt.show()

import pandas as pd
import numpy as np


def add_technical_indicators(data):
    """添加技术指标"""
    # 移动平均线
    data['MA5'] = data['close'].rolling(window=5).mean()
    data['MA20'] = data['close'].rolling(window=20).mean()

    # 相对强弱指数 (RSI)
    delta = data['close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    data['RSI'] = 100 - (100 / (1 + gain / loss))

    # 布林带
    data['BB_upper'] = data['close'].rolling(window=20).mean() + 2 * data['close'].rolling(window=20).std()
    data['BB_lower'] = data['close'].rolling(window=20).mean() - 2 * data['close'].rolling(window=20).std()

    # 去掉计算指标时产生的NaN值
    data = data.dropna()
    return data


# 添加技术指标
stock_a = add_technical_indicators(stock_a)
stock_b = add_technical_indicators(stock_b)

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error


def train_predict_model(data, target_col='close', features=None, days=5):
    """训练和预测模型"""
    if features is None:
        features = ['MA5', 'MA20', 'RSI', 'BB_upper', 'BB_lower']

    # 构建特征和目标
    X = []
    y = []
    for i in range(days, len(data)):
        X.append(data[features].iloc[i - days:i].values.flatten())
        y.append(data[target_col].iloc[i])

    X = np.array(X)
    y = np.array(y)

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 训练模型
    model = LinearRegression()
    model.fit(X_train, y_train)

    # 预测
    y_pred = model.predict(X_test)

    # 评估模型
    mse = mean_squared_error(y_test, y_pred)
    print(f"均方误差 (MSE): {mse}")

    return model, X_test, y_test, y_pred


# 训练模型
model_a, X_test_a, y_test_a, y_pred_a = train_predict_model(stock_a)
model_b, X_test_b, y_test_b, y_pred_b = train_predict_model(stock_b)

# 可视化预测结果
plt.figure(figsize=(12, 6))
plt.plot(y_test_a, label='实际价格')
plt.plot(y_pred_a, label='预测价格', linestyle='--')
plt.legend()
plt.title(' 中国平安价格预测')
plt.show()

plt.figure(figsize=(12, 6))
plt.plot(y_test_b, label='实际价格')
plt.plot(y_pred_b, label='预测价格', linestyle='--')
plt.legend()
plt.title(' 招商银行价格预测')
plt.show()


def ml_trading_strategy(data, model, features=None, days=5):
    """基于机器学习的交易策略"""
    if features is None:
        features = ['MA5', 'MA20', 'RSI', 'BB_upper', 'BB_lower']

    signals = pd.DataFrame(index=data.index)
    signals['price'] = data['close']
    signals['signal'] = 0.0

    for i in range(days, len(data)):
        X = data[features].iloc[i - days:i].values.flatten().reshape(1, -1)
        predicted_price = model.predict(X)
        if predicted_price > data['close'].iloc[i - 1]:
            signals.loc[data.index[i], 'signal'] = 1.0  # 买入
        elif predicted_price < data['close'].iloc[i - 1]:
            signals.loc[data.index[i], 'signal'] = -1.0  # 卖出

    signals['positions'] = signals['signal'].diff()
    return signals


# 应用交易策略
ml_signals_a = ml_trading_strategy(stock_a, model_a)
ml_signals_b = ml_trading_strategy(stock_b, model_b)

# 可视化交易信号
plt.figure(figsize=(12, 6))
plt.plot(ml_signals_a['price'], label='Price')
plt.plot(ml_signals_a['price'][ml_signals_a['positions'] == 1], '^', markersize=10, color='g', label='买入')
plt.plot(ml_signals_a['price'][ml_signals_a['positions'] == -1], 'v', markersize=10, color='r', label='卖出')
plt.legend()
plt.title('基于机器学习的交易策略 - 中国银行')
plt.show()

plt.figure(figsize=(12, 6))
plt.plot(ml_signals_b['price'], label='Price')
plt.plot(ml_signals_b['price'][ml_signals_b['positions'] == 1], '^', markersize=10, color='g', label='买入')
plt.plot(ml_signals_b['price'][ml_signals_b['positions'] == -1], 'v', markersize=10, color='r', label='卖出')
plt.legend()
plt.title('基于机器学习的交易策略 - 招商银行')
plt.show()

def ml_trading_strategy_with_sl_tp(data, model, features=None, days=5, stop_loss=0.05, take_profit=0.05):
    """基于机器学习的交易策略，包含止损和止盈"""
    if features is None:
        features = ['MA5', 'MA20', 'RSI', 'BB_upper', 'BB_lower']

    signals = pd.DataFrame(index=data.index)
    signals['price'] = data['close']
    signals['signal'] = 0.0
    signals['position'] = 0.0
    signals['stop_loss'] = 0.0
    signals['take_profit'] = 0.0

    for i in range(days, len(data)):
        X = data[features].iloc[i - days:i].values.flatten().reshape(1, -1)
        predicted_price = model.predict(X)
        current_price = data['close'].iloc[i - 1]

        if predicted_price > current_price:
            signals.loc[data.index[i], 'signal'] = 1.0  # 买入
        elif predicted_price < current_price:
            signals.loc[data.index[i], 'signal'] = -1.0  # 卖出

        if signals['position'].iloc[i - 1] == 1.0:  # 持有多头仓位
            if (current_price - signals['price'].iloc[i - 1]) / signals['price'].iloc[i - 1] <= -stop_loss:
                signals.loc[data.index[i], 'signal'] = -1.0  # 止损卖出
            elif (current_price - signals['price'].iloc[i - 1]) / signals['price'].iloc[i - 1] >= take_profit:
                signals.loc[data.index[i], 'signal'] = -1.0  # 止盈卖出
        elif signals['position'].iloc[i - 1] == -1.0:  # 持有空头仓位
            if (signals['price'].iloc[i - 1] - current_price) / signals['price'].iloc[i - 1] <= -stop_loss:
                signals.loc[data.index[i], 'signal'] = 1.0  # 止损买入
            elif (signals['price'].iloc[i - 1] - current_price) / signals['price'].iloc[i - 1] >= take_profit:
                signals.loc[data.index[i], 'signal'] = 1.0  # 止盈买入

        signals.loc[data.index[i], 'position'] = signals['position'].iloc[i - 1] + signals['signal'].iloc[i]

    signals['positions'] = signals['signal'].diff()
    return signals

def backtest_strategy(signals, initial_capital=1000000):
    """回测策略"""
    positions = signals['position']
    portfolio = pd.DataFrame(index=signals.index)
    portfolio['holdings'] = positions * signals['price']
    portfolio['cash'] = initial_capital - (positions.diff() * signals['price']).cumsum()
    portfolio['total'] = portfolio['cash'] + portfolio['holdings']
    portfolio['returns'] = portfolio['total'].pct_change()

    return portfolio

# 回测多支股票
def backtest_multiple_stocks(stocks, models, features, days=5, stop_loss=0.05, take_profit=0.05, initial_capital=1000000):
    """回测多支股票"""
    total_portfolio = pd.DataFrame(index=stocks[0].index)
    total_portfolio['total'] = 0.0

    for stock, model in zip(stocks, models):
        signals = ml_trading_strategy_with_sl_tp(stock, model, features, days, stop_loss, take_profit)
        portfolio = backtest_strategy(signals, initial_capital)
        total_portfolio['total'] += portfolio['total']

    total_portfolio['returns'] = total_portfolio['total'].pct_change()
    return total_portfolio

# 添加技术指标
stock_a = add_technical_indicators(stock_a)
stock_b = add_technical_indicators(stock_b)

# 训练模型
model_a, _, _, _ = train_predict_model(stock_a)
model_b, _, _, _ = train_predict_model(stock_b)

# 回测多支股票
stocks = [stock_a, stock_b]
models = [model_a, model_b]
features = ['MA5', 'MA20', 'RSI', 'BB_upper', 'BB_lower']
days = 5
stop_loss = 0.05
take_profit = 0.05
initial_capital = 1000000

total_portfolio = backtest_multiple_stocks(stocks, models, features, days, stop_loss, take_profit, initial_capital)

# 可视化回测结果
plt.figure(figsize=(12, 6))
plt.plot(total_portfolio['total'], label='总组合价值')
plt.legend()
plt.title('多支股票回测结果')
plt.show()

# 计算收益率
total_return = (total_portfolio['total'].iloc[-1] - initial_capital) / initial_capital
print(f"总收益率: {total_return:.2%}")

import numpy as np

def calculate_sharpe_ratio(returns, risk_free_rate=0.01):
    """计算夏普比率"""
    excess_returns = returns - risk_free_rate
    sharpe_ratio = np.mean(excess_returns) / np.std(excess_returns)
    return sharpe_ratio

# 计算夏普比率
sharpe_ratio = calculate_sharpe_ratio(total_portfolio['returns'])
print(f"夏普比率: {sharpe_ratio:.2f}")